import sys
import pyaudio
from PyQt5.QtWidgets import QApplication, QWidget, QVBoxLayout, QPushButton, QLabel, QLineEdit, QFileDialog
from PyQt5.QtCore import QTimer
import numpy as np
from PyQt5.QtGui import QPainter, QPen, QColor
from PyQt5.QtCore import Qt
import wave  # 用于读取WAV文件格式相关信息

class VoiceRecorder(QWidget):
    def __init__(self):
        super().__init__()

        self.initUI()
        self.initAudio()

    def initUI(self):
        self.setWindowTitle('语音信号采集与基本参数分析系统')

        self.layout = QVBoxLayout()

        # 状态标签
        self.status_label = QLabel('状态:准备好开始')
        self.layout.addWidget(self.status_label)

        # 开始录音按钮
        self.record_button = QPushButton('开始录音')
        self.record_button.clicked.connect(self.start_recording)
        self.layout.addWidget(self.record_button)

        # 停止录音按钮
        self.stop_button = QPushButton('停止录音')
        self.stop_button.clicked.connect(self.stop_recording)
        self.stop_button.setEnabled(False)
        self.layout.addWidget(self.stop_button)

        # 音频信息标签
        self.info_label = QLabel('语音信息:')
        self.layout.addWidget(self.info_label)

        # 添加采样率设置相关控件
        self.rate_label = QLabel('采样率 (Hz):')
        self.layout.addWidget(self.rate_label)
        self.rate_edit = QLineEdit(str(44100))  # 默认值设为44100
        self.layout.addWidget(self.rate_edit)

        # 添加声道数设置相关控件
        self.channels_label = QLabel('声道数:')
        self.layout.addWidget(self.channels_label)
        self.channels_edit = QLineEdit(str(1))  # 默认值设为1
        self.layout.addWidget(self.channels_edit)

        # 添加每块帧数大小设置相关控件
        self.chunk_label = QLabel('每块帧数大小:')
        self.layout.addWidget(self.chunk_label)
        self.chunk_edit = QLineEdit(str(1024))  # 默认值设为1024
        self.layout.addWidget(self.chunk_edit)

        # 添加保存按钮
        self.save_button = QPushButton('保存音频到本地')
        self.save_button.clicked.connect(self.save_audio)
        self.save_button.setEnabled(False)
        self.layout.addWidget(self.save_button)

        # 添加打开文件按钮
        self.open_button = QPushButton('打开文件')
        self.open_button.clicked.connect(self.open_audio_file)
        self.layout.addWidget(self.open_button)

        # 添加倍速变换相关控件
        self.speed_label = QLabel('倍速 (例如：2.0表示2倍速)：')
        self.layout.addWidget(self.speed_label)
        self.speed_edit = QLineEdit('1.0')  # 默认1倍速
        self.layout.addWidget(self.speed_edit)
        self.apply_speed_button = QPushButton('应用倍速')
        self.apply_speed_button.clicked.connect(self.apply_speed_change)
        self.layout.addWidget(self.apply_speed_button)

        self.setLayout(self.layout)

    def initAudio(self):
        self.p = pyaudio.PyAudio()
        self.stream = None
        self.frames = []

        # 初始音频参数，后续会根据输入控件的值更新
        self.format = pyaudio.paInt16
        self.channels = 1
        self.rate = 44100
        self.chunk = 1024

    def start_recording(self):
        # 获取并更新音频参数
        self.rate = int(self.rate_edit.text())
        self.channels = int(self.channels_edit.text())
        self.chunk = int(self.chunk_edit.text())

        self.status_label.setText('状态: 录音中...')
        self.stop_button.setEnabled(True)
        self.record_button.setEnabled(False)
        self.save_button.setEnabled(False)  # 开始录音时禁用保存按钮

        # 根据更新后的参数打开音频流
        self.stream = self.p.open(format=self.format,
                                  channels=self.channels,
                                  rate=self.rate,
                                  input=True,
                                  frames_per_buffer=self.chunk)

        # 启动定时器读取音频流
        self.timer = QTimer()
        self.timer.timeout.connect(self.read_audio_stream)
        self.timer.start(10)  # 10 milliseconds

    def read_audio_stream(self):
        data = self.stream.read(self.chunk)
        self.frames.append(data)

    def stop_recording(self):
        self.timer.stop()
        self.stream.stop_stream()
        self.stream.close()
        self.status_label.setText('Status: Ready to record')
        self.record_button.setEnabled(True)
        self.stop_button.setEnabled(False)
        self.save_button.setEnabled(True)


        self.analyze_audio()

    def analyze_audio(self):
        # 计算大致录音时间（秒）
        recording_time = len(self.frames) * self.chunk / self.rate

        # 转换音频帧为numpy数组用于绘图
        audio_data = b''.join(self.frames)
        audio_array = np.frombuffer(audio_data, dtype=np.int16)

        # 更新音频信息标签文本
        self.info_label.setText(f'声音信息:\n'
                                f'采样率: {self.rate} Hz\n'
                                f'声道: {self.channels}\n'
                                f'Format: 16-bit int\n'
                                f'总Frame: {len(self.frames)}\n'
                                f'记录时间: {recording_time:.2f} 秒')

        self.waveform_widget = WaveformWidget(audio_array, 800, 400)
        self.layout.addWidget(self.waveform_widget)
        self.update()

        with open('output.wav', 'wb') as wf:
            wf.write(b'RIFF')
            wf.write((len(audio_data) + 36).to_bytes(4, 'little'))
            wf.write(b'WAVEfmt ')
            wf.write((16).to_bytes(4, 'little'))
            wf.write((1).to_bytes(2, 'little'))
            wf.write(self.channels.to_bytes(2, 'little'))
            wf.write(self.rate.to_bytes(4, 'little'))
            wf.write((self.rate * self.channels * 2).to_bytes(4, 'little'))
            wf.write((self.channels * 2).to_bytes(2, 'little'))
            wf.write((16).to_bytes(2, 'little'))
            wf.write(b'data')
            wf.write(len(audio_data).to_bytes(4, 'little'))
            # 写入音频数据
            wf.write(audio_data)

    def save_audio(self):
        file_path, _ = QFileDialog.getSaveFileName(self, "保存音频文件", "", "WAV Files (*.wav)")
        if file_path:
            # 转wav
            recording_time = len(self.frames) * self.chunk / self.rate

            # 转换音频帧为numpy数组用于绘图
            audio_data = b''.join(self.frames)
            audio_array = np.frombuffer(audio_data, dtype=np.int16)
            with open(file_path, 'wb') as wf:
                # 写入WAV文件头
                wf.write(b'RIFF')
                wf.write((len(audio_data) + 36).to_bytes(4, 'little'))
                wf.write(b'WAVEfmt ')
                wf.write((16).to_bytes(4, 'little'))
                wf.write((1).to_bytes(2, 'little'))
                wf.write(self.channels.to_bytes(2, 'little'))
                wf.write(self.rate.to_bytes(4, 'little'))
                wf.write((self.rate * self.channels * 2).to_bytes(4, 'little'))
                wf.write((self.channels * 2).to_bytes(2, 'little'))
                wf.write((16).to_bytes(2, 'little'))
                wf.write(b'data')
                wf.write(len(audio_data).to_bytes(4, 'little'))
                wf.write(audio_data)

    def open_audio_file(self):
        file_path, _ = QFileDialog.getOpenFileName(self, "打开音频文件", "", "WAV Files (*.wav)")
        if file_path:
            try:
                with wave.open(file_path, 'rb') as wf:
                    # 获取音频文件参数
                    self.channels = wf.getnchannels()
                    self.rate = wf.getframerate()
                    self.format = pyaudio.paInt16
                    self.frames = []
                    n_frames = wf.getnframes()
                    self.chunk = 1024
                    data = wf.readframes(self.chunk)
                    while data:
                        self.frames.append(data)
                        data = wf.readframes(self.chunk)

                    # 转换音频帧为numpy数组用于绘图及显示信息
                    audio_data = b''.join(self.frames)
                    audio_array = np.frombuffer(audio_data, dtype=np.int16)


                    recording_time = n_frames / self.rate
                    self.info_label.setText(f'声音信息:\n'
                                            f'采样率: {self.rate} Hz\n'
                                            f'声道: {self.channels}\n'
                                            f'Format: 16-bit int\n'
                                            f'总Frame: {n_frames}\n'
                                            f'时间: {recording_time:.2f} 秒')


                    self.waveform_widget = WaveformWidget(audio_array, 800, 400)
                    self.layout.addWidget(self.waveform_widget)
                    self.update()
            except:
                print("打开文件出错，请确保文件格式正确且可正常读取")

    def apply_speed_change(self):
        try:
            speed = float(self.speed_edit.text())
            if speed <= 0:
                raise ValueError("倍速必须大于0")

            new_frames = []
            interval = int(1 / speed)
            for i in range(0, len(self.frames), interval):
                new_frames.append(self.frames[i])

            self.frames = new_frames


            audio_data = b''.join(self.frames)
            audio_array = np.frombuffer(audio_data, dtype=np.int16)


            recording_time = len(self.frames) * self.chunk / self.rate

            # 更新音频信息标签文本
            self.info_label.setText(f'Audio Info:\n'
                                    f'Sampling Rate: {self.rate} Hz\n'
                                    f'Channels: {self.channels}\n'
                                    f'Format: 16-bit int\n'
                                    f'Total Frames: {len(self.frames)}\n'
                                    f'倍速: {speed}倍\n'
                                    f'新时长: {recording_time:.2f} 秒')


            self.layout.removeWidget(self.waveform_widget)
            self.waveform_widget.deleteLater()
            self.waveform_widget = WaveformWidget(audio_array, 800, 400)
            self.layout.addWidget(self.waveform_widget)
            self.update()
        except ValueError as e:
            print(f"输入倍速有误: {e}")

    def closeEvent(self, event):
        if self.stream:
            self.stream.stop_stream()
            self.stream.close()
        self.p.terminate()
        event.accept()


class WaveformWidget(QWidget):
    def __init__(self, audio_data, width, height):
        super().__init__()
        self.audio_data = audio_data
        self.width_value = width
        self.height_value = height
        self.setMinimumSize(800, 400)  # 设置部件最小尺寸

    def paintEvent(self, event):
        painter = QPainter(self)
        painter.setPen(QPen(QColor(0, 0, 0), 1))  # 设置画笔颜色和宽度

        # 缩放音频数据以适配部件高度
        audio_data_scaled = self.audio_data / np.max(np.abs(self.audio_data)) * (self.height_value / 2)

        x_scale = self.width_value / len(self.audio_data)
        for i in range(len(self.audio_data) - 1):
            x1 = i * x_scale
            y1 = self.height_value / 2 - audio_data_scaled[i]
            x2 = (i + 1) * x_scale
            y2 = self.height_value / 2 - audio_data_scaled[i + 1]
            painter.drawLine(int(x1), int(y1), int(x2), int(y2))


if __name__ == '__main__':
    app = QApplication(sys.argv)
    ex = VoiceRecorder()
    ex.show()
    sys.exit(app.exec_())